import os
from hashlib import md5 as MD5
import shutil

try:
    from cvio import cvio
except Exception as e:
    print(e)
    from .cvio import cvio


def gethash(img):
    with open(img, 'rb') as fp:
        return MD5(fp.read()).hexdigest()

def gethashes(imgs):
    hashes = []
    n = len(imgs)
    for i, img in enumerate(imgs,1):
        x = gethash(img)
        print('calculate hash code', '[%d/%d]' % (i, n), x)
        hashes.append(x)
    return hashes

def findRepeatImgs(src1, src2, dst, with_json=True, recursive=True):
    images1 = cvio.load_image_list(src1, recursive=recursive)
    images2 = cvio.load_image_list(src2, recursive=recursive)

    hashes1 = gethashes(images1)
    hashes2 = gethashes(images2)

    repeat_hashes = list(set(hashes1).intersection(set(hashes2)))
    if len(repeat_hashes) < 1:
        print('Not found any repeated images in %s' % (src2))
        return 
    repeat_images = [images2[i] for i, x in enumerate(hashes2) if x in repeat_hashes]
    n = len(repeat_images)
    if not os.path.exists(dst):
        os.makedirs(dst)
    for i, img in enumerate(repeat_images,1):
        print('Moving [%d/%d] %s' % (i, n, img))
        shutil.move(img, os.path.join(dst, os.path.basename(img)))
        if with_json:
            js = os.path.splitext(img)[0] + '.json'
            shutil.move(js, os.path.join(dst, os.path.basename(js)))


if __name__ == '__main__':
    src1 = r'G:\data\datasets\drink\pesi\ai\pesi_test_200'
    src2 = r'G:\data\datasets\drink\pesi\ai\pesi_test_508'
    dst = r'G:\data\datasets\drink\pesi\ai\pesi_test_200_'
    findRepeatImgs(src1, src2, dst, recursive=False)